package d_scene.a_2_rate_limiter;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶算法
 *
 * @author He Kuntao
 * @since 2025/3/10 15:57
 */
public class TokenBucketRateLimiter {

    /**
     * 令牌桶容量
     */
    private final int capacity;
    /**
     * 每秒令牌生成速率
     */
    private final int refillRate;

    /**
     * 当前令牌数量
     */
    private final AtomicInteger tokens;


    public TokenBucketRateLimiter(int capacity, int refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = new AtomicInteger(0);
        startRefilling();
    }

    /**
     * 定期往令牌桶中添加令牌
     */
    private void startRefilling() {
        ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(1);
        scheduledExecutorService.scheduleWithFixedDelay(() -> {
            int currentTokens = tokens.get();
            if (currentTokens < capacity) {
                int newTokens = Math.min(capacity, currentTokens + refillRate);
                tokens.set(newTokens);
                System.out.println("新增令牌，当前令牌数：" + newTokens);
            }

        }, 0, 1, TimeUnit.SECONDS);
    }

    /**
     * 获取令牌
     *
     * @return true 表示获取令牌成功，false 表示获取令牌失败
     */
    public boolean tryAcquire() {
        int currentTokens = tokens.get();
        if (currentTokens > 0) {
            if (tokens.compareAndSet(currentTokens, currentTokens - 1)) {
                System.out.println("获取令牌成功");
                return true;
            } else {
                System.out.println("获取令牌失败 CAS False");
                return false;
            }
        }
        System.out.println("获取令牌失败 No Capacity");
        return false;
    }

    public static void main(String[] args) {
        // 初始容量为10，每秒生成2个令牌
        TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(10, 2);

        // 模拟10个请求
        for (int i = 0; i < 100; i++) {
            boolean acquire = rateLimiter.tryAcquire();
            if (!acquire) {
                System.out.printf("[%d请求] 被拒绝\n", i);
            }else {
                System.out.printf("[%d请求] 被执行\n", i);
            }
            try {
                Thread.sleep(200); // 模拟请求间隔
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }
}
